//! # JWT认证提供者
//!
//! 实现基于JSON Web Token的认证和授权功能，支持标准的JWT令牌验证和权限检查。
//!
//! ## 功能特性
//!
//! * **JWT令牌验证** - 支持HS256、RS256等签名算法
//! * **自动过期检查** - 自动验证令牌是否过期
//! * **Claims提取** - 从JWT中提取用户信息和权限
//! * **自定义Claims** - 支持自定义声明字段
//! * **令牌刷新** - 支持刷新令牌机制
//!
//! ## 使用示例
//!
//! ```rust
//! use rustcloud_gateway::auth::jwt::*;
//!
//! // 创建JWT认证提供者
//! let jwt_config = JwtConfig {
//!     secret: "your-secret-key".to_string(),
//!     algorithm: JwtAlgorithm::HS256,
//!     issuer: Some("rustcloud".to_string()),
//!     ..Default::default()
//! };
//! let provider = JwtAuthProvider::new(jwt_config).await?;
//!
//! // 认证请求
//! let context = AuthContext::new(headers, method, uri, remote_addr);
//! let principal = provider.authenticate(&context).await?;
//! ```

use super::{AuthProvider, AuthContext, Principal, Resource, AuthResult, AuthError, 
           AuthProviderType, PrincipalType, PermissionChecker};
use rustcloud_security::providers::AuthorizationProvider;
// 导入新的 Principal 构建器工具和权限检查器
use rustcloud_core::auth::{PrincipalConverter, PermissionCheckerFactory};
use async_trait::async_trait;
use jsonwebtoken::{decode, encode, Algorithm, DecodingKey, EncodingKey, Header, Validation};
use serde::{Deserialize, Serialize};
use std::collections::{HashMap, HashSet};
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use std::sync::Arc;
use tracing::{debug, error, warn};

/// JWT算法类型
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub enum JwtAlgorithm {
    /// HMAC SHA-256
    HS256,
    /// HMAC SHA-384
    HS384,
    /// HMAC SHA-512
    HS512,
    /// RSASSA-PKCS1-v1_5 SHA-256
    RS256,
    /// RSASSA-PKCS1-v1_5 SHA-384
    RS384,
    /// RSASSA-PKCS1-v1_5 SHA-512
    RS512,
}

impl From<JwtAlgorithm> for Algorithm {
    fn from(alg: JwtAlgorithm) -> Self {
        match alg {
            JwtAlgorithm::HS256 => Algorithm::HS256,
            JwtAlgorithm::HS384 => Algorithm::HS384,
            JwtAlgorithm::HS512 => Algorithm::HS512,
            JwtAlgorithm::RS256 => Algorithm::RS256,
            JwtAlgorithm::RS384 => Algorithm::RS384,
            JwtAlgorithm::RS512 => Algorithm::RS512,
        }
    }
}

/// JWT配置
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct JwtConfig {
    /// 签名密钥或公钥
    pub secret: String,
    /// 私钥（用于签名，可选）
    pub private_key: Option<String>,
    /// 算法类型
    pub algorithm: JwtAlgorithm,
    /// 发行者
    pub issuer: Option<String>,
    /// 受众
    pub audience: Option<Vec<String>>,
    /// 令牌有效期
    pub expiration: Duration,
    /// 允许的时钟偏差
    pub leeway: Duration,
    /// 是否验证过期时间
    pub validate_exp: bool,
    /// 是否验证生效时间
    pub validate_nbf: bool,
    /// 自定义声明字段映射
    pub claims_mapping: ClaimsMapping,
}

impl Default for JwtConfig {
    fn default() -> Self {
        Self {
            secret: "default-secret-key".to_string(),
            private_key: None,
            algorithm: JwtAlgorithm::HS256,
            issuer: Some("rustcloud-gateway".to_string()),
            audience: None,
            expiration: Duration::from_secs(24 * 60 * 60), // 24 hours
            leeway: Duration::from_secs(60),
            validate_exp: true,
            validate_nbf: true,
            claims_mapping: ClaimsMapping::default(),
        }
    }
}

/// Claims字段映射配置
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ClaimsMapping {
    /// 用户ID字段名
    pub user_id_field: String,
    /// 用户名字段名
    pub username_field: String,
    /// 角色字段名
    pub roles_field: String,
    /// 权限字段名
    pub permissions_field: String,
    /// 用户类型字段名
    pub user_type_field: String,
}

impl Default for ClaimsMapping {
    fn default() -> Self {
        Self {
            user_id_field: "sub".to_string(),
            username_field: "name".to_string(),
            roles_field: "roles".to_string(),
            permissions_field: "permissions".to_string(),
            user_type_field: "user_type".to_string(),
        }
    }
}

/// JWT声明
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct JwtClaims {
    /// 主题（用户ID）
    pub sub: String,
    /// 发行者
    pub iss: Option<String>,
    /// 受众
    pub aud: Option<Vec<String>>,
    /// 过期时间
    pub exp: i64,
    /// 生效时间
    pub nbf: Option<i64>,
    /// 签发时间
    pub iat: Option<i64>,
    /// JWT ID
    pub jti: Option<String>,
    /// 用户名
    pub name: Option<String>,
    /// 角色列表
    pub roles: Option<Vec<String>>,
    /// 权限列表
    pub permissions: Option<Vec<String>>,
    /// 用户类型
    pub user_type: Option<String>,
    /// 自定义声明
    #[serde(flatten)]
    pub custom_claims: HashMap<String, serde_json::Value>,
}

impl JwtClaims {
    /// 创建新的JWT声明
    pub fn new(user_id: String, expiration: Duration) -> Self {
        let now = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap()
            .as_secs() as i64;
        
        Self {
            sub: user_id,
            iss: None,
            aud: None,
            exp: now + expiration.as_secs() as i64,
            nbf: Some(now),
            iat: Some(now),
            jti: None,
            name: None,
            roles: None,
            permissions: None,
            user_type: None,
            custom_claims: HashMap::new(),
        }
    }

    /// 添加自定义声明
    pub fn add_custom_claim(mut self, key: String, value: serde_json::Value) -> Self {
        self.custom_claims.insert(key, value);
        self
    }

    /// 检查是否过期
    pub fn is_expired(&self) -> bool {
        let now = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap()
            .as_secs() as i64;
        now >= self.exp
    }

    /// 检查是否有效
    pub fn is_valid(&self) -> bool {
        let now = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap()
            .as_secs() as i64;
        
        // 检查生效时间
        if let Some(nbf) = self.nbf {
            if now < nbf {
                return false;
            }
        }
        
        // 检查过期时间
        !self.is_expired()
    }
}

/// JWT认证提供者
pub struct JwtAuthProvider {
    config: JwtConfig,
    encoding_key: Option<EncodingKey>,
    decoding_key: DecodingKey,
    validation: Validation,
    /// 权限检查器（使用新的统一实现）
    permission_checker: Arc<dyn PermissionChecker>,
}

impl JwtAuthProvider {
    /// 创建新的JWT认证提供者
    pub async fn new(config: JwtConfig) -> AuthResult<Self> {
        let algorithm: Algorithm = config.algorithm.clone().into();
        
        // 创建解码密钥
        let decoding_key = match &config.algorithm {
            JwtAlgorithm::HS256 | JwtAlgorithm::HS384 | JwtAlgorithm::HS512 => {
                DecodingKey::from_secret(config.secret.as_bytes())
            }
            JwtAlgorithm::RS256 | JwtAlgorithm::RS384 | JwtAlgorithm::RS512 => {
                DecodingKey::from_rsa_pem(config.secret.as_bytes())
                    .map_err(|e| AuthError::InternalError(format!("Invalid RSA public key: {}", e)))?
            }
        };

        // 创建编码密钥（如果需要）
        let encoding_key = if let Some(private_key) = &config.private_key {
            Some(match &config.algorithm {
                JwtAlgorithm::HS256 | JwtAlgorithm::HS384 | JwtAlgorithm::HS512 => {
                    EncodingKey::from_secret(private_key.as_bytes())
                }
                JwtAlgorithm::RS256 | JwtAlgorithm::RS384 | JwtAlgorithm::RS512 => {
                    EncodingKey::from_rsa_pem(private_key.as_bytes())
                        .map_err(|e| AuthError::InternalError(format!("Invalid RSA private key: {}", e)))?
                }
            })
        } else if matches!(config.algorithm, JwtAlgorithm::HS256 | JwtAlgorithm::HS384 | JwtAlgorithm::HS512) {
            Some(EncodingKey::from_secret(config.secret.as_bytes()))
        } else {
            None
        };

        // 创建验证配置
        let mut validation = Validation::new(algorithm);
        validation.validate_exp = config.validate_exp;
        validation.validate_nbf = config.validate_nbf;
        validation.leeway = config.leeway.as_secs();
        
        if let Some(ref issuer) = config.issuer {
            validation.iss = Some([issuer.clone()].into_iter().collect());
        }
        
        if let Some(ref audience) = config.audience {
            validation.aud = Some(audience.iter().cloned().collect());
        }

        Ok(Self {
            config,
            encoding_key,
            decoding_key,
            validation,
            // 使用新的统一权限检查器
            permission_checker: Arc::new(PermissionCheckerFactory::create_default()),
        })
    }

    /// 生成JWT令牌
    pub async fn generate_token(&self, claims: &JwtClaims) -> AuthResult<String> {
        let encoding_key = self.encoding_key.as_ref()
            .ok_or_else(|| AuthError::InternalError("No encoding key available".to_string()))?;

        let header = Header::new(self.config.algorithm.clone().into());
        
        encode(&header, claims, encoding_key)
            .map_err(|e| AuthError::InternalError(format!("Failed to encode JWT: {}", e)))
    }

    /// 解析JWT令牌
    pub async fn parse_token(&self, token: &str) -> AuthResult<JwtClaims> {
        decode::<JwtClaims>(token, &self.decoding_key, &self.validation)
            .map(|token_data| token_data.claims)
            .map_err(|e| {
                warn!("JWT parsing failed: {}", e);
                match e.kind() {
                    jsonwebtoken::errors::ErrorKind::ExpiredSignature => AuthError::TokenExpired,
                    jsonwebtoken::errors::ErrorKind::InvalidToken => AuthError::InvalidToken(e.to_string()),
                    _ => AuthError::AuthenticationFailed(format!("JWT validation failed: {}", e)),
                }
            })
    }

    /// 从Claims创建Principal（使用新的构建器工具）
    fn claims_to_principal(&self, claims: &JwtClaims) -> Principal {
        // 将 JWT Claims 转换为 HashMap 格式，以便使用 PrincipalConverter
        let mut claims_map = HashMap::new();
        
        // 添加基本信息
        if let Some(name) = &claims.name {
            claims_map.insert("name".to_string(), serde_json::Value::String(name.clone()));
        }
        
        if let Some(user_type) = &claims.user_type {
            claims_map.insert("user_type".to_string(), serde_json::Value::String(user_type.clone()));
        }
        
        // 添加角色
        if let Some(roles) = &claims.roles {
            let roles_array = roles.iter()
                .map(|r| serde_json::Value::String(r.clone()))
                .collect::<Vec<_>>();
            claims_map.insert("roles".to_string(), serde_json::Value::Array(roles_array));
        }
        
        // 添加权限
        if let Some(permissions) = &claims.permissions {
            let permissions_array = permissions.iter()
                .map(|p| serde_json::Value::String(p.clone()))
                .collect::<Vec<_>>();
            claims_map.insert("permissions".to_string(), serde_json::Value::Array(permissions_array));
        }
        
        // 添加过期时间
        claims_map.insert("exp".to_string(), serde_json::Value::Number(serde_json::Number::from(claims.exp)));
        
        // 添加自定义 Claims
        for (key, value) in &claims.custom_claims {
            claims_map.insert(key.clone(), value.clone());
        }
        
        // 使用 PrincipalConverter 创建 Principal
        PrincipalConverter::from_jwt_claims(claims.sub.clone(), &claims_map).build()
    }
}

#[async_trait]
impl AuthProvider for JwtAuthProvider {
    async fn authenticate(&self, context: &AuthContext) -> AuthResult<Principal> {
        debug!("Authenticating JWT token for request: {} {}", context.method, context.uri);

        // 获取Bearer令牌
        let token = context.get_bearer_token()
            .ok_or_else(|| AuthError::AuthenticationFailed("Missing Bearer token".to_string()))?;

        // 解析和验证令牌
        let claims = self.parse_token(&token).await?;

        // 检查令牌是否有效
        if !claims.is_valid() {
            return Err(AuthError::TokenExpired);
        }

        // 创建Principal
        let principal = self.claims_to_principal(&claims);

        debug!("JWT authentication successful for user: {}", principal.id);
        Ok(principal)
    }

    async fn validate_token(&self, token: &str) -> AuthResult<bool> {
        match self.parse_token(token).await {
            Ok(claims) => Ok(claims.is_valid()),
            Err(_) => Ok(false),
        }
    }

    fn provider_type(&self) -> AuthProviderType {
        AuthProviderType::Jwt
    }

    fn provider_name(&self) -> &str {
        "jwt"
    }

    fn supports_token_refresh(&self) -> bool {
        true
    }

    async fn refresh_token(&self, token: &str) -> AuthResult<String> {
        // 解析现有令牌
        let claims = self.parse_token(token).await?;
        
        // 创建新的过期时间
        let mut new_claims = claims;
        let now = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap()
            .as_secs() as i64;
        new_claims.exp = now + self.config.expiration.as_secs() as i64;
        new_claims.iat = Some(now);

        // 生成新令牌
        self.generate_token(&new_claims).await
    }

    async fn logout(&self, _principal: &Principal) -> AuthResult<()> {
        // JWT 通常是无状态的，不需要特殊的注销处理
        // 但可以在这里添加令牌黑名单逻辑
        Ok(())
    }

    async fn get_user_info(&self, principal: &Principal) -> AuthResult<serde_json::Value> {
        Ok(serde_json::json!({
            "id": principal.id,
            "name": principal.name,
            "type": principal.principal_type,
            "roles": principal.roles,
            "permissions": principal.permissions,
            "attributes": principal.attributes
        }))
    }
}

/// JWT授权提供者实现（使用新的统一权限检查器）
#[async_trait]
impl AuthorizationProvider for JwtAuthProvider {
    async fn authorize(&self, principal: &Principal, resource: &Resource) -> AuthResult<bool> {
        debug!("Authorizing user {} for resource {}", principal.id, resource.get_resource_id());

        // 使用新的统一权限检查器
        self.permission_checker.check_resource(principal, resource).await
    }

    async fn check_permission(&self, principal: &Principal, permission: &str) -> AuthResult<bool> {
        // 使用新的统一权限检查器
        self.permission_checker.check(principal, permission).await
    }

    async fn get_user_permissions(&self, principal: &Principal) -> AuthResult<Vec<String>> {
        Ok(principal.get_permissions().iter().cloned().collect())
    }

    async fn get_user_roles(&self, principal: &Principal) -> AuthResult<Vec<String>> {
        Ok(principal.get_roles().iter().cloned().collect())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use axum::http::HeaderMap;
    use std::time::Duration;

    #[tokio::test]
    async fn test_jwt_provider_creation() {
        let config = JwtConfig::default();
        let provider = JwtAuthProvider::new(config).await;
        assert!(provider.is_ok());
    }

    #[tokio::test]
    async fn test_jwt_token_generation() {
        let config = JwtConfig::default();
        let provider = JwtAuthProvider::new(config).await.unwrap();

        let claims = JwtClaims::new("user123".to_string(), Duration::from_secs(3600)) // 1 hour
            .add_custom_claim("name".to_string(), serde_json::Value::String("Test User".to_string()));

        let token = provider.generate_token(&claims).await;
        assert!(token.is_ok());
    }

    #[tokio::test]
    async fn test_jwt_authentication_success() {
        let config = JwtConfig::default();
        let provider = JwtAuthProvider::new(config).await.unwrap();

        // 生成测试令牌
        let claims = JwtClaims::new("user123".to_string(), Duration::from_secs(3600)); // 1 hour
        let token = provider.generate_token(&claims).await.unwrap();

        // 创建认证上下文
        let mut headers = HeaderMap::new();
        headers.insert("authorization", format!("Bearer {}", token).parse().unwrap());
        
        let context = AuthContext::new(
            headers,
            "GET".to_string(),
            "/api/test".to_string(),
            Some("127.0.0.1".to_string()),
        );

        // 执行认证
        let result = provider.authenticate(&context).await;
        assert!(result.is_ok());
        
        let principal = result.unwrap();
        assert_eq!(principal.id, "user123");
    }

    #[tokio::test]
    async fn test_jwt_authentication_missing_token() {
        let config = JwtConfig::default();
        let provider = JwtAuthProvider::new(config).await.unwrap();

        let headers = HeaderMap::new();
        let context = AuthContext::new(
            headers,
            "GET".to_string(),
            "/api/test".to_string(),
            Some("127.0.0.1".to_string()),
        );

        let result = provider.authenticate(&context).await;
        assert!(result.is_err());
        assert!(matches!(result.unwrap_err(), AuthError::AuthenticationFailed(_)));
    }

    #[tokio::test]
    async fn test_jwt_authorization_admin_role() {
        let config = JwtConfig::default();
        let provider = JwtAuthProvider::new(config).await.unwrap();

        let principal = Principal::new(
            "admin1".to_string(),
            "Admin User".to_string(),
            PrincipalType::Admin,
        ).add_role("admin".to_string());

        let resource = Resource::new(
            "user-service".to_string(),
            "/api/users".to_string(),
            "GET".to_string(),
        );

        // 使用 AuthorizationProvider trait
        let auth_provider: &dyn AuthorizationProvider = &provider;
        let result = auth_provider.authorize(&principal, &resource).await;
        assert!(result.is_ok());
        assert!(result.unwrap());
    }

    #[tokio::test]
    async fn test_jwt_authorization_specific_permission() {
        let config = JwtConfig::default();
        let provider = JwtAuthProvider::new(config).await.unwrap();

        let principal = Principal::new(
            "user1".to_string(),
            "Regular User".to_string(),
            PrincipalType::User,
        ).add_permission("user-service:get:/api/users".to_string());

        let resource = Resource::new(
            "user-service".to_string(),
            "/api/users".to_string(),
            "GET".to_string(),
        );

        // 使用 AuthorizationProvider trait
        let auth_provider: &dyn AuthorizationProvider = &provider;
        let result = auth_provider.authorize(&principal, &resource).await;
        assert!(result.is_ok());
        assert!(result.unwrap());
    }

    #[tokio::test]
    async fn test_jwt_authorization_denied() {
        let config = JwtConfig::default();
        let provider = JwtAuthProvider::new(config).await.unwrap();

        let principal = Principal::new(
            "user1".to_string(),
            "Regular User".to_string(),
            PrincipalType::User,
        );

        let resource = Resource::new(
            "user-service".to_string(),
            "/api/admin".to_string(),
            "DELETE".to_string(),
        );

        // 使用 AuthorizationProvider trait
        let auth_provider: &dyn AuthorizationProvider = &provider;
        let result = auth_provider.authorize(&principal, &resource).await;
        assert!(result.is_ok());
        assert!(!result.unwrap());
    }

    #[tokio::test]
    async fn test_jwt_token_validation() {
        let config = JwtConfig::default();
        let provider = JwtAuthProvider::new(config).await.unwrap();

        let claims = JwtClaims::new("user123".to_string(), Duration::from_secs(3600)); // 1 hour
        let token = provider.generate_token(&claims).await.unwrap();

        let is_valid = provider.validate_token(&token).await.unwrap();
        assert!(is_valid);

        // 测试无效令牌
        let is_valid = provider.validate_token("invalid.token.here").await.unwrap();
        assert!(!is_valid);
    }

    #[tokio::test]
    async fn test_jwt_token_refresh() {
        let config = JwtConfig::default();
        let provider = JwtAuthProvider::new(config).await.unwrap();

        let claims = JwtClaims::new("user123".to_string(), Duration::from_secs(3600)); // 1 hour
        let original_token = provider.generate_token(&claims).await.unwrap();

        let new_token = provider.refresh_token(&original_token).await.unwrap();
        assert_ne!(original_token, new_token);
        
        let is_valid = provider.validate_token(&new_token).await.unwrap();
        assert!(is_valid);
    }
}